package mcpplus

import (
	"context"
	"errors"
	"fmt"
	"gitee.com/hongzhaomin/gspring/annotation"
	"gitee.com/hongzhaomin/gspring/aware"
	"gitee.com/hongzhaomin/gspring/iface"
	"gitee.com/hongzhaomin/gspringboot"
	"gitee.com/hongzhaomin/hzm-common-go/assert"
	"github.com/mark3labs/mcp-go/server"
	"net/http"
	"reflect"
	"strings"
)

var _ aware.BeanFactoryAware = (*McpServerManager)(nil)

type McpServerManager struct {
	annotation.Component
	aware.Aware
	annotation.Initialize

	beanFactory iface.BeanFactory
	logger      iface.GspringLogger
	environment gspringboot.ConfigurationEnvironment

	mcpServer *server.MCPServer
	sseServer *server.SSEServer
}

func (my *McpServerManager) SetBeanFactory(factory iface.BeanFactory) {
	my.beanFactory = factory
	my.logger = factory.GetBeanByRt(reflect.TypeOf((*iface.GspringLogger)(nil))).(iface.GspringLogger)
	my.environment = factory.GetBeanByRt(reflect.TypeOf((*gspringboot.ConfigurationEnvironment)(nil))).(gspringboot.ConfigurationEnvironment)
}

func (my *McpServerManager) Init() {
	mcpServer := my.beanFactory.GetBeans(reflect.TypeOf((*server.MCPServer)(nil)))
	if len(mcpServer) > 0 {
		my.mcpServer = mcpServer[0].(*server.MCPServer)
	} else {
		mcpName := my.environment.GetStringReplacePlaceholders(gspringAiMcpNamePropName)
		assert.NotBlank(mcpName, fmt.Sprintf("property [%s] is blank", gspringAiMcpNamePropName))
		mcpVersion := my.environment.GetStringReplacePlaceholders(gspringAiMcpVersionPropName)
		assert.NotBlank(mcpVersion, fmt.Sprintf("property [%s] is blank", gspringAiMcpVersionPropName))
		// Create a new MCP server
		my.mcpServer = server.NewMCPServer(
			mcpName,
			mcpVersion,
			server.WithResourceCapabilities(true, true),
			server.WithLogging(),
			server.WithRecovery(),
		)
	}
}

func (my *McpServerManager) Start() {
	mcpMode := my.environment.GetStringReplacePlaceholders(gspringAiMcpModePropName)
	assert.NotBlank(mcpMode, fmt.Sprintf("property [%s] is blank", gspringAiMcpModePropName))

	my.registerConfiguration()

	// 异步启动mcp服务
	switch strings.ToLower(mcpMode) {
	case mcpModeStdio:
		go func() {
			if err := server.ServeStdio(my.mcpServer); err != nil {
				my.logger.Error("failed to start mcp server for stdio: %v", err)
			}
		}()
	case mcpModeSSE:
		mcpPort := "8080"
		my.setPropIfExist(gspringAiMcpPortPropName, func(val string) {
			mcpPort = val
		})
		go func() {
			my.sseServer = server.NewSSEServer(my.mcpServer)
			if err := my.sseServer.Start(":" + mcpPort); err != nil {
				if errors.Is(err, http.ErrServerClosed) {
					my.logger.Info("mcp server closed")
				} else {
					my.logger.Error("failed to start mcp server for sse: %v", err)
				}
			}
		}()
	default:
		panic(errors.New(fmt.Sprintf("mcp mode %s is not supported", mcpMode)))
	}
}

func (my *McpServerManager) Shutdown() {
	if my.sseServer != nil {
		if err := my.sseServer.Shutdown(context.Background()); err != nil {
			my.logger.Error("failed to shutdown mcp server for sse: %v", err)
		}
	}
}

func (my *McpServerManager) setPropIfExist(key string, set func(val string)) {
	environment := my.environment
	if environment.ContainsKey(key) {
		set(environment.GetStringReplacePlaceholders(key))
	}
}

func (my *McpServerManager) registerConfiguration() {
	logger := my.logger
	factory := my.beanFactory
	// 注册 [ToolsAnnotation]
	tools := factory.GetBeans(reflect.TypeOf((*ToolsAnnotation)(nil)))
	toolsBean := make([]ToolsAnnotation, len(tools))
	for i, tool := range tools {
		toolsBean[i] = tool.(ToolsAnnotation)
	}

	parser := NewToolParser()
	parser.ParseTools(toolsBean...)
	for _, mcpTool := range parser.mcpTools {
		my.mcpServer.AddTool(mcpTool, parser.ToolHandlerFunc)
		logger.Info("The mcp tool have completed registration -> %s", mcpTool.Name)
	}
}
